#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
数据读取模块
专门负责Excel文件的读取和基础数据处理
"""

import pandas as pd
import numpy as np
from pathlib import Path
import warnings
warnings.filterwarnings('ignore')

class EpidemicDataReader:
    """疫情数据读取器"""
    
    def __init__(self, file_path="香港各区疫情数据_20250322.xlsx"):
        """
        初始化数据读取器
        
        Args:
            file_path (str): Excel文件路径
        """
        self.file_path = Path(file_path)
        self.data = None
        self.sheets_info = {}
        
    def load_data(self):
        """加载Excel数据"""
        try:
            if not self.file_path.exists():
                raise FileNotFoundError(f"文件不存在: {self.file_path}")
            
            # 读取Excel文件的所有工作表
            excel_file = pd.ExcelFile(self.file_path)
            self.sheets_info = {
                'sheet_names': excel_file.sheet_names,
                'total_sheets': len(excel_file.sheet_names)
            }
            
            print(f"✅ 成功加载Excel文件: {self.file_path}")
            print(f"📊 工作表数量: {self.sheets_info['total_sheets']}")
            print(f"📋 工作表名称: {', '.join(self.sheets_info['sheet_names'])}")
            
            return True
            
        except Exception as e:
            print(f"❌ 加载数据失败: {str(e)}")
            return False
    
    def read_sheet(self, sheet_name=None, sheet_index=0):
        """
        读取指定的工作表
        
        Args:
            sheet_name (str): 工作表名称
            sheet_index (int): 工作表索引
        """
        try:
            if sheet_name:
                self.data = pd.read_excel(self.file_path, sheet_name=sheet_name)
            else:
                # 使用sheet_name参数，通过索引获取工作表名称
                excel_file = pd.ExcelFile(self.file_path)
                sheet_name_by_index = excel_file.sheet_names[sheet_index]
                self.data = pd.read_excel(self.file_path, sheet_name=sheet_name_by_index)
            
            print(f"✅ 成功读取工作表: {sheet_name if sheet_name else f'索引{sheet_index}'}")
            print(f"📈 数据形状: {self.data.shape}")
            print(f"📋 列名: {list(self.data.columns)}")
            
            return True
            
        except Exception as e:
            print(f"❌ 读取工作表失败: {str(e)}")
            return False
    
    def get_data_info(self):
        """获取数据基本信息"""
        if self.data is None:
            print("❌ 请先加载数据")
            return
        
        print("\n" + "="*50)
        print("📊 数据基本信息")
        print("="*50)
        
        # 基本信息
        print(f"数据形状: {self.data.shape}")
        print(f"行数: {self.data.shape[0]}")
        print(f"列数: {self.data.shape[1]}")
        
        # 数据类型
        print(f"\n数据类型:")
        print(self.data.dtypes)
        
        # 缺失值信息
        print(f"\n缺失值统计:")
        missing_data = self.data.isnull().sum()
        if missing_data.sum() > 0:
            print(missing_data[missing_data > 0])
        else:
            print("✅ 无缺失值")
        
        # 数据预览
        print(f"\n前20行数据:")
        print(self.data.head(20))
        
        # 数值列统计
        numeric_cols = self.data.select_dtypes(include=[np.number]).columns
        if len(numeric_cols) > 0:
            print(f"\n数值列统计信息:")
            print(self.data[numeric_cols].describe())
    
    def get_sheets_summary(self):
        """获取所有工作表摘要"""
        if not self.sheets_info:
            print("❌ 请先加载Excel文件")
            return
        
        print("\n" + "="*50)
        print("📋 工作表摘要")
        print("="*50)
        
        for i, sheet_name in enumerate(self.sheets_info['sheet_names']):
            try:
                sheet_data = pd.read_excel(self.file_path, sheet_name=sheet_name)
                print(f"\n工作表 {i+1}: {sheet_name}")
                print(f"  形状: {sheet_data.shape}")
                print(f"  列名: {list(sheet_data.columns)}")
                
                # 显示前几行
                print(f"  前3行数据:")
                print(sheet_data.head(3).to_string())
                
            except Exception as e:
                print(f"❌ 读取工作表 {sheet_name} 失败: {str(e)}")
    
    def get_data(self):
        """获取当前加载的数据"""
        return self.data
    
    def get_sheets_info(self):
        """获取工作表信息"""
        return self.sheets_info 